import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.getcwd(), ".")))
from collections import OrderedDict
from typing import Dict, Iterable, Callable

import torch
from torch import nn, Tensor
import timm
import torch.nn as nn



class IntermediateLayerGetter(nn.ModuleDict):
    """
    Module wrapper that returns intermediate layers from a model

    It has a strong assumption that the modules have been registered
    into the model in the same order as they are used.
    This means that one should **not** reuse the same nn.Module
    twice in the forward if you want this to work.

    Additionally, it is only able to query submodules that are directly
    assigned to the model. So if `model` is passed, `model.feature1` can
    be returned, but not `model.feature1.layer2`.

    Args:
        model (nn.Module): model on which we will extract the features
        return_layers (Dict[name, new_name]): a dict containing the names
            of the modules for which the activations will be returned as
            the key of the dict, and the value of the dict is the name
            of the returned activation (which the user can specify).

    Examples::

        >>> m = torchvision.models.resnet18(weights=ResNet18_Weights.DEFAULT)
        >>> # extract layer1 and layer3, giving as names `feat1` and feat2`
        >>> new_m = torchvision.models._utils.IntermediateLayerGetter(m,
        >>>     {'layer1': 'feat1', 'layer3': 'feat2'})
        >>> out = new_m(torch.rand(1, 3, 224, 224))
        >>> print([(k, v.shape) for k, v in out.items()])
        >>>     [('feat1', torch.Size([1, 64, 56, 56])),
        >>>      ('feat2', torch.Size([1, 256, 14, 14]))]
    """

    _version = 2
    __annotations__ = {
        "return_layers": Dict[str, str],
    }

    def __init__(self, model: nn.Module, return_layers: Dict[str, str]) -> None:
        if not set(return_layers).issubset([name for name, _ in model.named_children()]):
            raise ValueError("return_layers are not present in model")
        orig_return_layers = return_layers
        return_layers = {str(k): str(v) for k, v in return_layers.items()}
        layers = OrderedDict()
        for name, module in model.named_children():
            layers[name] = module
            if name in return_layers:
                del return_layers[name]
            if not return_layers:
                break

        super().__init__(layers)
        self.return_layers = orig_return_layers

    def forward(self, x):
        out = OrderedDict()
        for name, module in self.items():
            x = module(x)
            if name in self.return_layers:
                out_name = self.return_layers[name]
                out[out_name] = x
        return out


class VerboseExecution(nn.Module):
    def __init__(self, model: nn.Module):
        super().__init__()
        self.model = model

        # Register a hook for each layer
        for name, layer in self.model.named_children():
            layer.__name__ = name
            layer.register_forward_hook(
                lambda layer, _, output: print(f"{layer.__name__}: {output.shape}")
            )

    def forward(self, x: Tensor) -> Tensor:
        return self.model(x)


class FeatureExtractor(nn.Module):
    def __init__(self, model: nn.Module, layers: Iterable[str]):
        super().__init__()
        self.model = model
        self.layers = layers
        self._features = OrderedDict({layer: torch.empty(0) for layer in layers})
        self.hook = []

        for layer_id in layers:
            layer = dict([*self.model.named_modules()])[layer_id]
            self.layer_id = layer.register_forward_hook(self.save_outputs_hook(layer_id))
            self.hook.append(self.layer_id)

    def save_outputs_hook(self, layer_id: str) -> Callable:
        def fn(_, __, output):
            self._features[layer_id] = output

        return fn

    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        _ = self.model(x)
        return self._features

    def remove(self):
        for hook in self.hook:
            hook.remove()


class TimmEncoder(nn.Module):
    def __init__(self, name, pretrained=True, in_channels=3, features_only=True, depth=4, output_stride=32):
        super().__init__()
        kwargs = dict(
            in_chans=in_channels,
            pretrained=pretrained,
            features_only=features_only,
            output_stride=output_stride,
            out_indices=tuple(range(4)),
        )
        self.features_only = features_only

        if not self.features_only:
            kwargs.pop("output_stride")
            kwargs.pop("out_indices")

        elif output_stride == 32:
            # not all models support output stride argument, drop it by default
            kwargs.pop("output_stride")
        else:
            pass

        self.model = timm.create_model(name, **kwargs)

        self._in_channels = in_channels

        self._out_channels = self.model.feature_info.channels() if self.features_only else \
        list(self.model.feature_info)[-1]["num_chs"]
        self._depth = depth
        self._output_stride = output_stride

    def forward(self, x):
        if self.features_only:
            features = self.model(x)
        else:
            features = self.model.forward_features(x)

        return features

    @property
    def out_channels(self):
        return self._out_channels

    @property
    def output_stride(self):
        return min(self._output_stride, 2 ** self._depth)


if __name__ == "__main__":
    # , output_stride =16, out_indices=tuple(range(4))
    model = timm.create_model(model_name="resnet50",
                              in_chans=3,
                              features_only=True,
                              pretrained=False,
                              #   output_stride =16,
                              out_indices=tuple(range(1, 1 + 4)))
    # last = [list(model.feature_info)[-1]["num_chs"]]
    # print(last)
    inputs = torch.randn(2, 3, 256, 256)
    outs = model(inputs)
    # outs = model.forward_features(inputs)
    # print(outs.shape)
    print([o.shape for o in outs])
    # print(model)
    # avail_pretrained_models = timm.list_models(pretrained=True)
    # print(avail_pretrained_models)
